Closing in favor of #6878, which predates this PR and covers the same #6877 fix more completely, including the schema-less IterableDatasetDict case.

Your diagnosis is right and the one-liner is the correct fix. I checked out the branch and ran it both ways rather than reading the diff, and along the way found a second copy of the same bug twenty lines down that this PR does not fix.

Confirming yours first. On main, list(unpair_preference_dataset(idd)["t"]) raises IndexError: list index out of range, which is worth putting in the PR text since it is the string someone will search for. The cause is what you say: IterableDatasetDict.column_names is {'t': ['prompt', 'chosen', 'rejected']}, and that dict is truthy, so the or list(next(iter(dataset)).keys()) fallback never runs and remove_columns receives it. With your change the same call returns 4 rows from 2 prompts with completion, label, prompt. Reverting only the data_utils.py hunk and keeping your test gives 1 failed, 9 passed, and the one failure is your new test with that same IndexError. Your widened isinstance also matches what _get_dataset_format already does at data_utils.py:671.

The second copy is maybe_unpair_preference_dataset at data_utils.py:547:

if isinstance(dataset, DatasetDict):
    column_names = dataset[list(dataset.keys())[0]].column_names
else:
    column_names = dataset.column_names        # IterableDatasetDict -> the same dict
if "chosen" in column_names and "rejected" in column_names:

For an IterableDatasetDict that membership test looks at the dict's keys, which are split names, so it is False. Run on your branch:

'chosen' in column_names -> False
result rows = [{'prompt': 'a', 'chosen': 'x', 'rejected': 'p'},
               {'prompt': 'b', 'chosen': 'y', 'rejected': 'q'}]
n_rows = 2   (unchanged; 4 would mean unpaired)

It returns the paired dataset and says nothing, which is a worse failure than the one you fixed because there is no signal at all. The function is public, exported at trl/__init__.py:43, but the only in-tree caller is the experimental BCO trainer, which does not look like it accepts iterable datasets, so today you would reach it by calling the function directly.

Both functions sit next to each other and share one root cause, so I would fix them together here rather than in a second PR. Your call, and I would not block on it.

The test builds its expected rows by zipping to_dict().values(). Four literal dicts would be easier to check by eye, and would not quietly follow along if unpaired_dataset ever changes.

Thanks for catching something that only shows up on the streaming path.